import json
import numpy as np
import tensorrt as trt
from argparse import ArgumentParser
from skimage import io
from skimage.transform import resize
from pycuda import driver as cuda
import pycuda.autoinit
from preprocess import preprocess, simple_preprocess

with open('class_id.json') as f:
    class_id = json.load(f)

parser = ArgumentParser(description='run tensorrt resnet on image')
parser.add_argument('images', type=str, nargs='+')
ARGS = parser.parse_args()
images = ARGS.images

logger = trt.Logger(trt.Logger.INFO)
with open('resnet18_softmax_engine.trt', 'rb') as f:
    with trt.Runtime(logger) as runtime:
        engine = runtime.deserialize_cuda_engine(f.read())
        
        num_layers = engine.num_layers
        print('num layers:', num_layers)
        inspector = engine.create_engine_inspector() 
        for i in range(num_layers):
            print('layer #', i)
            info = inspector.get_layer_information(i, trt.LayerInformationFormat(1))
            print(info)


        for binding in engine:
            if engine.binding_is_input(binding):
                input_binding = binding
                input_shape = engine.get_binding_shape(binding)
                input_size = trt.volume(input_shape) * engine.max_batch_size * np.dtype(np.float32).itemsize
                print('binding:', binding, 'is input, shape:', input_shape)
            else:
                output_binding = binding
                output_shape = engine.get_binding_shape(binding)
                output_size = trt.volume(output_shape) * engine.max_batch_size * np.dtype(np.float32).itemsize
                print('binding:', binding, 'is output, shape:', output_shape)


context = engine.create_execution_context()
input_binding_index = engine.get_binding_index(input_binding)
output_binding_index = engine.get_binding_index(output_binding)
context.set_binding_shape(input_binding_index, input_shape)
stream = cuda.Stream()

device_input = cuda.mem_alloc(input_size)
host_output = cuda.pagelocked_empty(trt.volume(output_shape) * engine.max_batch_size, dtype=np.float32)
device_output = cuda.mem_alloc(host_output.nbytes)

for filename in images:
    image = io.imread(filename)
    print('image.shape:', image.shape, 'image.dtype:', image.dtype)
    host_input = simple_preprocess(image)
    print('host_input.shape:', host_input.shape)

    cuda.memcpy_htod_async(device_input, host_input, stream)
    context.execute_async_v2(bindings=[int(device_input), int(device_output)], stream_handle=stream.handle)
    cuda.memcpy_dtoh_async(host_output, device_output, stream)
    stream.synchronize()

    print('host_output.shape:', host_output.shape)
    class_index = np.argmax(host_output)
    confidence = host_output[class_index]
    print('filename:',  filename, 'class index:', class_index, 'class name:', class_id[class_index], 'confidence:', confidence)



